library(magrittr)
library(tidyverse)
library(Seurat)
library(readxl)
library(cowplot)
library(colorblindr)
library(viridis)
library(magick, lib.loc = "/home/uhlitzf/miniconda3/lib/R/library")
library(ggpubr)

1 Processing

## load global vars: 
source("_src/global_vars.R")

# meta_tbl
# clrs
# markers_v6
# markers_v6_super
# cell_type_super_lookup

names(clrs$cell_type) <- str_replace_all(names(clrs$cell_type), "\\.", " ")
names(clrs$cell_type) <- str_replace_all(names(clrs$cell_type), "Ovarian", "Ov")

## load data --------------------------------------

## load cohort embeddings
seu_tbl <- read_tsv("/work/shah/uhlitzf/data/SPECTRUM/freeze/v6/outs/cells.tsv") %>% 
  mutate(cell_type = ifelse(cell_type == "Monocyte", "Myeloid.cell", cell_type)) %>% 
  mutate(cell_type = ifelse(cell_type == "Ovarian.cancer.cell", "Ov.cancer.cell", cell_type))

seu_tbl <- read_tsv("/work/shah/isabl_data_lake/analyses/68/74/6874/cells.tsv") %>%
  mutate(cell_type = ifelse(cell_type == "Monocyte", "Myeloid.cell", cell_type)) %>%
  mutate(cell_type = ifelse(cell_type == "Ovarian.cancer.cell", "Ov.cancer.cell", cell_type))

## load consensus data
consOV_tbl <- read_tsv("/work/shah/uhlitzf/data/SPECTRUM/freeze/v6/consensusOV/SPECTRUM_freeze_v6_consensusOV.tsv") %>%
  mutate(consensusOV = ordered(consensusOV, levels = names(clrs$consensusOV))) %>% 
  select(cell_id, consensusOV)

## join data
seu_tbl_full <- seu_tbl %>% 
  left_join(meta_tbl, by = "sample") %>% 
  filter(therapy == "pre-Rx", cell_type != "Other", tumor_supersite != "Unknown") %>% 
  rename(UMAP_1 = umap50_1, UMAP_2 = umap50_2) %>% 
  mutate(cell_type_super = cell_type_super_lookup[cell_type]) %>% 
  mutate(cell_type = ordered(str_replace_all(cell_type, "\\.", " "), 
                             levels = names(clrs$cell_type))) %>% 
  mutate(sort_short_x = ifelse(sort_short == "U" & cell_type_super == "Immune", 
                               "CD45+", ifelse(sort_short == "U" & cell_type_super == "Stromal", 
                                               "CD45-", sort_short))) %>% 
  left_join(consOV_tbl, by = "cell_id")

# seu_tbl_full <- seu_tbl_full %>%
#   sample_n(10000)

1.1 UMAPs

dpi <- 72

base_umap <- ggplot(seu_tbl_full) +
  coord_fixed() +
  NoAxes() +
  theme(legend.position = c(0, 1),
        legend.justification = c("left", "top"),
        legend.box.just = "left",
        legend.margin = ggplot2::margin(1, 1, 1, 1),
        #panel.border = element_rect(linetype = 1, color = "black", size = 1),
        legend.text = element_text(size = 14, margin = ggplot2::margin(0, 10, 0, 0)),
        legend.spacing.x = unit(0, "npc"), 
        legend.title = element_blank(),
        plot.title = element_text(hjust = 0.5, vjust = 0.5, face = "plain", size = 22))

pt.size <- 0.1
pt.size.mini <- 0.01
pt.alpha <- 0.05
pt.alpha.mini <- 0.02

# median_tbl <- seu_tbl_full %>% 
#   filter(cell_type != "Other") %>% 
#   group_by(cell_type, reduction) %>% 
#   summarise(UMAP_1 = median(UMAP_1),
#             UMAP_2 = median(UMAP_2))

ncol1_legend <- list(
  guides(color = guide_legend(override.aes = list(size = 4, alpha = 1),
                              ncol = 1, 
                              label.position = "right")),
  theme(legend.text = element_text(size = 14, margin = ggplot2::margin(0, 10, 9.5, 0), color = "white"),
        legend.spacing.x = unit(0, "npc"))
)

umap_site <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = tumor_supersite), size = pt.size.mini, alpha = pt.alpha.mini, raster.dpi = dpi) +
  scale_color_manual(values = clrs$tumor_supersite) +
  guides(color = guide_legend(override.aes = list(size = 4, alpha = 1),
                              ncol = 2, 
                              label.position = "right"))


umap_site_mini <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = tumor_supersite), size = pt.size.mini, alpha = pt.alpha.mini, raster.dpi = dpi) +
  scale_color_manual(values = clrs$tumor_supersite) +
  guides(color = F)

umap_site_legend_ncol1 <- cowplot::get_legend(umap_site + ncol1_legend)
umap_site_legend <- cowplot::get_legend(umap_site)

umap_cell_type <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = fct_rev(cell_type)), size = pt.size, alpha = pt.alpha, raster.dpi = dpi) +
  scale_color_manual(values = clrs$cell_type) +
  ncol1_legend

umap_cell_type_mini <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = cell_type), size = pt.size.mini, alpha = pt.alpha.mini, raster.dpi = dpi) +
  scale_color_manual(values = clrs$cell_type) +
  guides(color = F)

umap_cell_type_legend <- cowplot::get_legend(umap_cell_type)

umap_patient <- base_umap + 
  geom_point(aes(UMAP_1, UMAP_2, color = patient_id_short), size = pt.size, alpha = pt.alpha) +
  scale_color_manual(values = clrs$patient_id_short) + 
  guides(color = guide_legend(override.aes = list(size = 4, alpha = 1),
                              nrow = 16, 
                              label.position = "right")) +
  theme(legend.text = element_text(size = 14, margin = ggplot2::margin(0, 10, 0, 0)),
        legend.spacing.x = unit(0, "npc"))

umap_patient_mini <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = patient_id_short), size = pt.size.mini, alpha = pt.alpha.mini, raster.dpi = dpi) +
  scale_color_manual(values = clrs$patient_id_short) +
  guides(color = F)

# umap_patient_legend <- cowplot::get_legend(umap_patient)

umap_consOV <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = consensusOV), size = pt.size, alpha = pt.alpha, raster.dpi = dpi) +
  scale_color_manual(values = clrs$consensusOV, labels = c("Immuno.", "Mesench.", "Prolif.", "Differen.")) + 
  guides(color = guide_legend(override.aes = list(size = 4, alpha = 1),
                              ncol = 2, 
                              label.position = "right")) +
  theme(legend.text = element_text(size = 14, margin = ggplot2::margin(0, 10, 0, 0)),
        legend.spacing.x = unit(0, "npc"))

umap_consOV_mini <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = consensusOV), size = pt.size.mini, alpha = pt.alpha.mini, raster.dpi = dpi) +
  scale_color_manual(values = clrs$consensusOV) +
  guides(color = F)

umap_consOV_legend <- cowplot::get_legend(umap_consOV)

umap_mutsig <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = consensus_signature), size = pt.size, alpha = pt.alpha, raster.dpi = dpi) +
  scale_color_manual(values = clrs$consensus_signature) + 
  guides(color = guide_legend(override.aes = list(size = 4, alpha = 1),
                              ncol = 2, 
                              label.position = "right")) +
  theme(legend.text = element_text(size = 14, margin = ggplot2::margin(0, 10, 0, 0)),
        legend.spacing.x = unit(0, "npc"))

umap_mutsig_mini <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = consensus_signature), size = pt.size.mini, alpha = pt.alpha.mini, raster.dpi = dpi) +
  scale_color_manual(values = clrs$consensus_signature) +
  guides(color = F)

umap_mutsig_legend <- cowplot::get_legend(umap_mutsig)

continuous_umap_layers <- list(
  scale_color_viridis_c(),
  guides(color = guide_colorbar(label.position = "right",
                                title.position = "top",
                                title.hjust = 0,
                                title.vjust = 1,
                                direction = "vertical")),
  theme(legend.title = element_text(face = "bold", size = 14),
        legend.key.height = unit(0.05, "npc"),
        legend.key.width = unit(0.03, "npc"),
        legend.position = c(-0.07, 1),
        legend.justification = c("left", "top"),
        plot.title = element_text(hjust = 0.5, vjust = 0.5, face = "plain", size = 18)))

umap_umis <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = log2(nCount_RNA)), size = pt.size, alpha = pt.alpha, raster.dpi = dpi) +
  labs(title = "log2(#UMIs)", color = "") +
  continuous_umap_layers

umap_genes <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = log2(nFeature_RNA)), size = pt.size, alpha = pt.alpha, raster.dpi = dpi) +
  labs(title = "log2(#Genes)", color = "") +
  continuous_umap_layers

umap_mitos <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = percent.mt), size = pt.size, alpha = pt.alpha, raster.dpi = dpi) +
  labs(title = "MT reads [%]", color = "") +
  continuous_umap_layers

umap_doublet_score <- base_umap + 
  ggrastr::geom_point_rast(aes(UMAP_1, UMAP_2, color = doublet_score), size = pt.size, alpha = pt.alpha, raster.dpi = dpi) +
  labs(title = "Doublet score", color = "") +
  continuous_umap_layers

1.2 total numbers

total_numbers_plot <- function(column_var) {
  column_var <- enquo(column_var)
  plot_data <- seu_tbl_full %>% 
    filter(cell_type != "Other", therapy == "pre-Rx") %>%
    group_by(!!column_var) %>%
    tally %>%
    arrange(n) %>%
    mutate(column_label = paste0(!!column_var, " (", format(n, trim=T, big.mark=","), ")")) %>% 
    mutate(!!column_var := ordered(!!column_var, levels = unique(!!column_var)))
  
  common_layers <- list(
    scale_y_continuous(expand = c(0, 0)),
    guides(fill = F, color = F),
    coord_flip(clip = "off"),
    theme(axis.title = element_blank(),
          axis.text = element_blank(),
          axis.line = element_blank(),
          axis.ticks = element_blank(),
          plot.title = element_text(face = "plain"))
  )
  
  p1 <- ggplot(plot_data) +
    geom_point(aes(!!column_var, 0, color = !!column_var), size = 4) +
    scale_color_manual(values = clrs[[as_label(column_var)]]) +
    common_layers 
    
  p2 <- ggplot(plot_data) +
    geom_bar(aes(!!column_var, n, fill = !!column_var), stat = "identity", width = 0.3) +
    geom_text(aes(!!column_var, 1, label = column_label), hjust = 0, vjust = 0, nudge_x = 0.2) +
    scale_fill_manual(values = clrs[[as_label(column_var)]]) +
    common_layers 
  
  plot_grid(p1, p2, align = "v", ncol = 2, rel_widths = c(0.03, 0.97))
  
}

numbers_plot_cell_type <- total_numbers_plot(cell_type)
numbers_plot_tumor_supersite <- total_numbers_plot(tumor_supersite)

1.3 Mixing per cell type

pm_score_tbl <- read_tsv("/work/shah/uhlitzf/data/SPECTRUM/freeze/v5/patient_mixing_scores.tsv") %>%
  mutate(cell_type = str_replace_all(cell_type, "Monocyte", "Myeloid cell")) %>% 
  mutate(cell_type = str_replace_all(cell_type, "Ovarian.cancer.cell", "Ov cancer cell"))

pm_score_tbl_ordered <- pm_score_tbl %>%
  filter(cell_type != "Other") %>%
  group_by(cell_type) %>%
  mutate(median = median(score)) %>%
  ungroup() %>%
  arrange(-median) %>%
  mutate(cell_type = ordered(cell_type, levels = unique(cell_type)))

mixing_plot_cell_type <- ggplot(pm_score_tbl_ordered) +
  geom_violin(aes(cell_type, score, fill = cell_type),
              adjust = 3, color = "white", width = 1) +
  geom_boxplot(aes(cell_type, score, color = cell_type),
               width = 0.2, size = 1, fill = NA, outlier.size = 0.01) +
  geom_boxplot(aes(cell_type, score, fill = cell_type),
               width = 0.2, color = "white", outlier.shape = NA) +
  geom_hline(yintercept = 0, linetype = 2, size = 0.5, color = "black") +
  scale_fill_manual(values = clrs$cell_type) +
  scale_color_manual(values = clrs$cell_type) +
  guides(fill = F, color = F) +
  scale_y_continuous(expand = c(0, 0),
                     limits = c(-5, 7),
                     breaks = c(-4, -2, 0, 2, 4, 6)) +
  labs(y = "Patient\nspecificity") +
  theme(axis.text.x = element_text(angle = 60, hjust = 1, vjust = 1),
        axis.title.x = element_blank())

# ggsave("_fig/_paper/01_patient_mixing_plot_v5.pdf", mixing_plot_cell_type,
#        width = 3, height = 5)
### pairwise comparisons

pm_score_summary <- pm_score_tbl_ordered %>%
  separate(cell_id, into = c("patient_id", "sample_suffix"), remove = F, extra = "drop", sep = "_") %>%
  group_by(patient_id, cell_type) %>%
  summarise(mean_score = mean(score))

ptt_tbl <- pairwise.t.test(pm_score_summary$mean_score, pm_score_summary$cell_type, p.adjust.method = "fdr") %$%
  p.value %>%
  as_tibble(rownames = "cell_type_1") %>%
  gather(cell_type_2, pvalue, -cell_type_1) %>%
  mutate(cell_type_2 = ordered(cell_type_2, levels = levels(pm_score_tbl_ordered$cell_type))) %>%
  mutate(cell_type_1 = ordered(cell_type_1, levels = levels(pm_score_tbl_ordered$cell_type))) %>%
  mutate(log10p = -log10(pvalue),
         sig = pvalue < 0.1,
         label = ifelse(sig, "*", "")) %>%
  na.omit()

pm_comparisons <- ggplot(filter(ptt_tbl, sig)) +
  # geom_tile(aes(fct_rev(cell_type_1), fct_rev(cell_type_2), fill = log10p)) +
  geom_point(aes(fct_rev(cell_type_1), fct_rev(cell_type_2), size = log10p), color = "black", shape = 21) +
  geom_point(aes(fct_rev(cell_type_1), fct_rev(cell_type_2), fill = log10p, size = log10p), shape = 21) +
  # geom_point(aes(fct_rev(cell_type_1), fct_rev(cell_type_2), shape = sig),
  #           color = "white", size = 5) +
  # geom_point(aes(fct_rev(cell_type_1), fct_rev(cell_type_2), shape = sig),
  #           color = "black", size = 3) +
  scale_fill_gradientn(colours = viridis(9)) +
  scale_size(guide = F) +
  # scale_color_gradientn(colours = viridis(9)) +
  theme(aspect.ratio = 1,
        axis.text.x = element_text(angle = 90, hjust = 1, vjust = 0.5),
        axis.title = element_blank(),
        plot.title = element_text(face = "plain", size = 12)) +
  scale_shape_manual(values = c("", "*"), guide = F) +
  # scale_x_discrete(expand = c(0, 0)) +
  # scale_y_discrete(expand = c(0, 0)) +
  labs(fill = "-log10(p-value)", title = "Pairwise comparisons\nof patient specificity")

1.4 consensusOV summary

consOV_tbl_summary <- seu_tbl_full %>%
  select(cell_type, consensusOV, patient_id_short, tumor_supersite) %>%
  mutate(consOV_cell_type = cell_type) %>%
  # mutate(consOV_cell_type = ifelse(S.Score > 0.5, "Cycling cell", as.character(cell_type))) %>%
  group_by(consOV_cell_type, consensusOV) %>%
  tally %>%
  group_by(consOV_cell_type) %>%
  mutate(nrel = n/(sum(n))*100) %>%
  ungroup %>%
  mutate(consensusOV = ordered(consensusOV, levels = c(names(clrs$consensusOV)))) %>%
  arrange(consensusOV, nrel)

heat_layers <- list(
  geom_tile(),
  scale_fill_gradientn(colors = magma(9)[c(-1)], limits = c(0, 100), breaks = c(0, 50, 100)),
  theme(
    axis.title = element_blank(),
    axis.ticks = element_blank(),
    axis.line = element_blank(),
    axis.text.x = element_text(angle = 60, hjust = 1, vjust = 1),
    legend.position = "bottom",
    aspect.ratio = 4/9
    ),
  guides(fill = guide_colorbar(title.position = "top"))
)

tcga_heat <- consOV_tbl_summary %>%
  filter(consOV_cell_type != "Other") %>%
  mutate(cell_type = ordered(consOV_cell_type, levels = unique(c("Myeloid cell", "Fibroblast", "Cycling cell", names(clrs$cell_type))))) %>%
  mutate(consensusOV = ordered(consensusOV, levels = names(clrs$consensusOV))) %>%
  ggplot(aes(cell_type, consensusOV, fill = nrel)) +
  labs(fill = "Fraction\nof cells [%]") +
  heat_layers

tcga_heat_wide <- consOV_tbl_summary %>%
  filter(consOV_cell_type != "Other") %>%
  mutate(cell_type = ordered(consOV_cell_type, levels = unique(c("Myeloid cell", "Fibroblast", "Cycling cell", names(clrs$cell_type))))) %>%
  mutate(consensusOV = ordered(consensusOV, levels = names(clrs$consensusOV))) %>%
  ggplot(aes(consensusOV, cell_type, fill = nrel)) +
  labs(fill = "Fraction\nof cells [%]") +
  heat_layers +
  theme(aspect.ratio = 9/4, legend.position = "right")

# tcga_heat + guides(fill = F)
tcga_heat_legend <- get_legend(tcga_heat)

2 Figure 2 upper panel

arrow <- arrow(angle = 20, type = "closed", length = unit(0.1, "npc"))
umap_coord_anno <- ggplot(tibble(group = c("UMAP1", "UMAP2"),
                                 x = c(0, 0), xend = c(1, 0),
                                 y = c(0, 0), yend = c(0, 1),
                                 lx = c(0.5, -0.15), ly = c(-0.15, 0.5),
                                 angle = c(0, 90))) +
  geom_segment(aes(x, y, xend = xend, yend = yend, group = group),
               arrow = arrow, size = 1, lineend = "round") +
  geom_text(aes(lx, ly, label = group, angle = angle), size = 4) +
  theme_void() +
  coord_fixed(xlim = c(-0.3, 1), ylim = c(-0.3, 1))

add_umap_coord <- function(gg_obj) {
  p <- ggdraw() + 
    draw_plot(gg_obj, x = 0, y = 0, width = 1, height = 1) +
    draw_plot(umap_coord_anno, x = -0.015, y = -0.02, width = 0.2, height = 0.2)
  return(p)
}

ggdraw() +
  ## left
  draw_label("Patient", x = 0.25, y = 0.98, size = 20) +
  draw_plot(add_umap_coord(umap_patient), x = 0, y = 0, width = 0.5, height = 1) +
  draw_label(paste0("n = ", nrow(seu_tbl_full)), 
             x = 0.25, y = 0.03, size = 16, hjust = 0) + 
  ## right upper
  draw_plot(mixing_plot_cell_type, x = 0.5+0.5*0, y = 0.66, width = 0.5*0.3, height = 0.33) +
  draw_label("Cell type", x = 0.71, y = 0.98, size = 20, hjust = 0) +
  draw_plot(umap_cell_type_mini, x = 0.83, y = 0.68, width = 0.5*1/3, height = 0.31) +
  draw_plot(numbers_plot_cell_type, x = 0.7, y = 0.651, width = 0.13, height = 0.315) +
  ## right middle
  draw_label("TCGA", x = 0.53+0.5*0, y = 0.6, size = 20, hjust = 0) +
  draw_plot(umap_consOV_mini, x = 0.5+0.5*0, y = 0.3, width = 0.5*1/3, height = 0.31) +
  draw_label("Signature", x = 0.53+0.5*0.33, y = 0.6, size = 20, hjust = 0) +
  draw_plot(umap_mutsig_mini, x = 0.5+0.5*0.33, y = 0.3, width = 0.5*1/3, height = 0.31) +
  draw_label("Site", x = 0.53+0.5*0.66, y = 0.6, size = 20, hjust = 0) + 
  draw_plot(umap_site_mini, x = 0.5+0.5*0.66, y = 0.3, width = 0.5*1/3, height = 0.31) +
  ## right lower
  draw_grob(umap_consOV_legend, x = 0.52+0.5*0, y = -0.67, hjust = 0, vjust = 0) +
  draw_plot(tcga_heat + guides(fill = F), x = 0.49+0.5*0, y = -0.03, width = 0.18, height = 0.3) +
  draw_grob(tcga_heat_legend, x = 0.68, y = -0.35) +
  draw_grob(umap_mutsig_legend, x = 0.52+0.5*0.33, y = -0.67, hjust = 0, vjust = 0) +
  draw_grob(umap_site_legend, x = 0.52+0.5*0.66, y = -0.67, hjust = 0, vjust = 0)

ggsave(filename = "_fig/002_cohort/002_cohort_umaps_v6.png", width = 20, height = 10)
ggsave(filename = "_fig/002_cohort/002_cohort_umaps_v6.pdf", width = 20, height = 10)
ggdraw() +
  ## left
  draw_label("Patient", x = 0.25, y = 0.98, size = 20) +
  draw_plot(add_umap_coord(umap_patient), x = 0, y = 0, width = 0.5, height = 1) +
  draw_label(paste0("n = ", nrow(seu_tbl_full)), 
             x = 0.25, y = 0.03, size = 16, hjust = 0) + 
  ## right upper
  draw_plot(mixing_plot_cell_type, x = 0.5+0.5*0, y = 0.66, width = 0.5*0.3, height = 0.33) +
  draw_label("Cell type", x = 0.725, y = 0.98, size = 20, hjust = 0) +
  draw_plot(umap_cell_type_mini, x = 0.83, y = 0.68, width = 0.5*1/3, height = 0.31) +
  draw_plot(numbers_plot_cell_type, x = 0.715, y = 0.651, width = 0.13, height = 0.315) +
  ## right middle
  draw_label("TCGA", x = 0.53+0.5*0, y = 0.6, size = 20, hjust = 0) +
  draw_plot(umap_consOV_mini, x = 0.5+0.5*0, y = 0.3, width = 0.5*1/3, height = 0.31) +
  draw_label("Site", x = 0.725, y = 0.6, size = 20, hjust = 0) +
  draw_plot(numbers_plot_tumor_supersite, x = 0.715, y = 0.34, width = 0.13, height = 0.25) +
  draw_plot(umap_site_mini, x = 0.5+0.5*0.66, y = 0.3, width = 0.5*1/3, height = 0.31) +
  ## right lower
  draw_plot(tcga_heat + guides(fill = F), x = 0.49+0.5*0, y = -0.03, width = 0.18, height = 0.3) +
  draw_grob(tcga_heat_legend, x = 0.68, y = -0.35) +
  draw_grob(umap_consOV_legend, x = 0.52+0.5*0, y = -0.67, hjust = 0, vjust = 0)

ggsave(filename = "_fig/002_cohort/002_cohort_umaps_v6.png", width = 20, height = 10)
ggsave(filename = "_fig/002_cohort/002_cohort_umaps_v6.pdf", width = 20, height = 10)

3 S1

## QC histograms
qc_tbl <- seu_tbl_full %>%
  mutate(`log2(#UMIs)` = log2(nCount_RNA),
         `log2(#Genes)` = log2(nFeature_RNA),
         `log2(#UMIs/#Genes)` = `log2(#UMIs)`-`log2(#Genes)`,
         `MT reads [%]` = percent.mt) %>%
  select(`log2(#UMIs)`,
         `log2(#Genes)`,
         # `log2(#UMIs/#Genes)`,
         `MT reads [%]`, cell_type) %>%
  gather(key, value, -cell_type) %>%
  group_by(key) %>%
  mutate(median_value = median(value))

qc_tbl <- bind_rows(qc_tbl, mutate(qc_tbl, cell_type = "All cell types")) %>%
  mutate(cell_type = ordered(cell_type,
                             levels = rev(c("All cell types", names(clrs$cell_type)))))

qc_plot <- ggplot(qc_tbl) +
  geom_violin(aes(cell_type, value, fill = cell_type),
              color = "white", width = 1) +
  geom_boxplot(aes(cell_type, value, color = cell_type),
               width = 0.15, size = 1, fill = NA, outlier.shape = NA) +
  geom_boxplot(aes(cell_type, value), fill = NA,
               width = 0.15, color = "white", outlier.shape = NA) +
  geom_hline(aes(yintercept = median_value),
             data = distinct(qc_tbl, key, median_value)) +
  facet_wrap(~key, scales = "free_x", strip.position = "bottom") +
  scale_fill_manual(values = c(clrs$cell_type, `All cell types` = "grey10"),
                    guide = F) +
  scale_color_manual(values = c(clrs$cell_type, `All cell types` = "grey10"),
                     guide = F) +
  coord_flip() +
  theme(axis.title = element_blank(),
        strip.placement = "outside")
custom_guide <- guide_legend(override.aes = list(size = 4, alpha = 1), 
                             ncol = 1, label.position = "right")

mutsig_legend <- cowplot::get_legend(umap_mutsig + theme(legend.title=element_text()) + labs(color = "Mutational\nsignature") + guides(color = custom_guide))
consOV_legend <- cowplot::get_legend(umap_consOV + theme(legend.title=element_text()) + labs(color = "TCGA subtype") + guides(color = custom_guide))
cell_type_legend <- cowplot::get_legend(umap_cell_type + theme(legend.title=element_text(), legend.text = element_text(size = 14, margin = ggplot2::margin(0, 10, 0, 0), color = "black")) + labs(color = "Cell type") + guides(color = custom_guide))
site_legend <- cowplot::get_legend(umap_site + theme(legend.title=element_text()) + labs(color = "Site") + guides(color = custom_guide))


figS1_panel <- ggdraw() +
  # draw_plot(add_umap_coord(umap_mutsig + guides(color = F) + labs(title = "Mutational signature")), width = 0.24, height = 0.48, x = 0.01, y = 0.5) +
  draw_plot(add_umap_coord(umap_umis), width = 0.24, height = 0.48, x = 0.01, y = 0.5) +
  draw_plot(add_umap_coord(umap_genes), width = 0.24, height = 0.48, x = 0.25, y = 0.5) +
  draw_plot(add_umap_coord(umap_mitos), width = 0.24, height = 0.48, x = 0.01, y = 0) +
  draw_plot(qc_plot, width = 0.5, height = 0.48, x = 0.5, y = 0.5) +
  draw_grob(mutsig_legend, x = 0.55, y = 0.48, hjust = 0, vjust = 1) +
  draw_grob(consOV_legend, x = 0.55, y = 0.25, hjust = 0, vjust = 1) +
  draw_grob(cell_type_legend, x = 0.7, y = 0.48, hjust = 0, vjust = 1) +
  draw_grob(site_legend, x = 0.85, y = 0.48, hjust = 0, vjust = 1) +
  draw_label("A", x = 0.01, y = 0.98, fontface = "bold", size = 22) +
  draw_label("B", x = 0.52, y = 0.98, fontface = "bold", size = 22) +
  draw_label("C", x = 0.52, y = 0.48, fontface = "bold", size = 22)

figS1_panel

ggsave("_fig/002_cohort/002_supplementary_umap.pdf", figS1_panel, width = 16, height = 9)
# xf <- 1.33
# 
# ggdraw() +
#   draw_plot(add_umap_coord(umap_site), x = 0, y = 0, width = 0.5/xf, height = 1) +
#   draw_plot(umap_patient_mini, x = 0.5/xf, y = 0.5, width = 0.25/xf, height = 0.5) +
#   draw_plot(umap_cell_type_mini, x = 0.5/xf, y = 0, width = 0.25/xf, height = 0.5) +
#   draw_plot(umap_mutsig_mini, x = 0.94/xf, y = 0.5, width = 0.25/xf, height = 0.5) +
#   draw_plot(umap_consOV_mini, x = 0.94/xf, y = 0, width = 0.25/xf, height = 0.5) +
#   draw_grob(umap_cell_type_legend, x = 0.56, y = -0.53, hjust = 0, vjust = 0) +
#   draw_plot(total_numbers_cell_type, x = 0.565, y = 0.005, width = 0.13, height = 0.45) +
#   draw_grob(umap_patient_legend, x = 0.56, y = -0.05, hjust = 0, vjust = 0) +
#   draw_grob(umap_mutsig_legend, x = 0.9, y = -0.05, hjust = 0, vjust = 0) +
#   draw_grob(umap_consOV_legend, x = 0.9, y = -0.55, hjust = 0, vjust = 0) +
#   draw_label("Site", x = 0.26/xf, y = 0.98, size = 20) +
#   draw_label("Patient", x = 0.52/xf, y = 0.98, size = 20, hjust = 0) +
#   draw_label("Cell type", x = 0.52/xf, y = 0.49, size = 20, hjust = 0) +
#   draw_label("Signature", x = 0.95/xf, y = 0.98, size = 20, hjust = 0) +
#   draw_label("TCGA", x = 0.95/xf, y = 0.49, size = 20, hjust = 0)
# 
# # ggsave(filename = "_fig/001x_umaps_mutsig.png", width = 16, height = 7)
# ggsave(filename = "_fig/001x_umaps_mutsig.pdf", width = 18, height = 7)